import os
from dotenv import load_dotenv
from langchain.chains import GraphCypherQAChain
from langchain_community.graphs import Neo4jGraph
from langchain.prompts import ChatPromptTemplate
from model.my_chat_model import ChatModel
import warnings
from langchain_core._api import LangChainDeprecationWarning

warnings.filterwarnings("ignore", category=LangChainDeprecationWarning)


class RagNeo4jService:

    # 数据初始化
    def __init__(self):
        # 1 加载环境变量
        load_dotenv()
        # 2 初始化聊天模型
        self.chat = ChatModel()
        self.llm = self.chat.get_line_model()
        # 3 初始化图数据库
        self.graph = Neo4jGraph(url=os.getenv("NEO4J_URI"), username=os.getenv("NEO4J_USER"),
                                password=os.getenv("NEO4J_PASSWORD"))

    # 创建一个对话链
    def create_chain(self):
        # 2 获取提示模板
        prompt = ChatPromptTemplate.from_messages([
            ("system", """你是一个专业的cypher语言的查询器，能根据用户的提问，完成基于cypher语言的查询,
            数据库模式：
            - 节点属性: 
                1. 企业节点：上市代码、信用评级、企业名称、报告期、主营业务、成立日期、注册资本、上市状态、统一社会信用代码、市盈率
                2. 高管节点：风险标签、报告期、name、所属部门、任职时间、高管姓名、是否兼职、职位
                3. 十大股东节点：变动趋势、股本性质、报告期、持股比例、股东名称、name、股东排名、出生年月、性别、持股数量
                4. 最终受益人节点：管辖区域、控制级别、受益人ID、名称、受益类型、实体类型、受益比例
            - 关系: 
                1. 持股：用于连接十大股东与企业（表示股东对企业的持股关系）
                2. 母公司持股：用于连接企业与企业（表示母公司对下属企业的持股关系）
                3. 任职于：用于连接高管与企业（表示高管在企业的任职关系）
            """),
            ("human", "{question}")
        ])
        # 创建链
        chain = GraphCypherQAChain.from_llm(
            llm=self.llm,  # llm模型
            graph=self.graph,  # 图形数据库链接
            allow_dangerous_requests=True,  # 允许执行Cypyhe语言
            cypher_prompt=prompt,  # 提示模板
            verbose=True,  # 输出详细日志,调试用的，可以省略
            top_k=10,  # 返回结果数量,类似于 return limit 10
            validate_cypher=True  # 验证cypher语言
        )
        return chain
